import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,tydat_NH,average_datetime,split_str_id
from geopy.distance import geodesic



names,track_ids = ['2023NH42_36','2023NH47_46'],[42,47]
tynames,tyids = split_str_id(names)


for num,(ty,tyid) in enumerate(zip(tynames,tyids)):
    ''' init '''
    RIstd = 7
    dt = timedelta(days=2)
    directory =  f'D:\\met_data\\ty_ensemble\\{names[num]}\\'    # 要筛选的文件开头
    dates_name = os.listdir(directory)
    dates_paths = [directory + date for date in dates_name]
    obs_path = r'D:\met_data\ty_obs\storms\storms_NH2023'
    pic_savepath = rf'D:\met_data\excels'
    os.makedirs(pic_savepath,exist_ok=True)
    
    
    ''' calc 2d numbers '''
    # 遍历所有起报时间
    max_delta_time = timedelta(days=0)
    for j,i in enumerate(dates_paths):
        mem_paths = glob.glob( os.path.join(i,"TRACK_ID_*") )
        # sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
        start_time = datetime.strptime(dates_name[j],"%Y%m%d%H")
        # 遍历所有成员
        for single_mem_path in mem_paths:
            try:
                single_mem = tydat(single_mem_path,RIstd) 
                delta_time = single_mem.time[-1] - start_time
                if delta_time>max_delta_time:
                    max_delta_time = delta_time
            except Exception as e:
                pass
        max_dt_num = math.ceil(max_delta_time/dt) 
    
    ''' 求各起报时间的图 '''
    tyobs = tydat_NH(obs_path,track_ids[num])
    rcd_obs = count_rapidgrow(RIstd, tyobs.umax, tyobs.time)
    num_result = np.zeros(( max_dt_num,len(dates_name) ))
    dis_result = np.zeros(( max_dt_num,len(dates_name) ))
    dt_result  = np.zeros(( max_dt_num,len(dates_name) ))
    
    
    for i in tqdm( range(len(dates_paths)), desc=f'processing charts {ty}'):
        '''读取这个起报时间的所有集合'''
        mem_paths = glob.glob( os.path.join(dates_paths[i],"TRACK_ID_*") )
        sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
        ty_mem = [tydat(i,RIstd) for i in sorted_mem_paths]
        start_dt = datetime.strptime( dates_name[i],"%Y%m%d%H" )
        for j in range(max_dt_num):
            t0 = 2*j
            t1 = 2*j+2
            t0_dt = start_dt + timedelta(days=t0)
            t1_dt = start_dt + timedelta(days=t1)
            '''如果观测在这两天没有RI，就直接跳过'''
            idx_obs_2d = (tyobs.time>=t0_dt)&(tyobs.time<t1_dt)
            idx_obs_2d_RI = (rcd_obs[idx_obs_2d]==1)
            if np.sum(rcd_obs[idx_obs_2d]) ==0.:
                ''' only record numbers here '''
                for k in range(len(ty_mem)):
                    '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                    idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                    idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                    num_result[j,i] += np.sum(idx_mem_2d_RI)
                continue
            else:
                '''如果2d内有多个RI,取平均'''
                time_obs_ave = average_datetime(tyobs.time[idx_obs_2d][idx_obs_2d_RI])
                lon_obs_ave  = np.mean(tyobs.lon[idx_obs_2d][idx_obs_2d_RI])
                lat_obs_ave  = np.mean(tyobs.lat[idx_obs_2d][idx_obs_2d_RI])
                '''开始计算'''
                num = 0 
                dis_sum = 0 
                dt_sum  = 0 
                for k in range(len(ty_mem)):
                    '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                    idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                    idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                    time_mem = ty_mem[k].time[idx_mem_2d][idx_mem_2d_RI]
                    lat_mem  = ty_mem[k].lat[idx_mem_2d][idx_mem_2d_RI]
                    lon_mem  = ty_mem[k].lon[idx_mem_2d][idx_mem_2d_RI]
                    if np.sum(idx_mem_2d_RI) ==0.:
                        continue
                    else:
                        num+=np.sum(idx_mem_2d_RI)
                        for l in range(len(lon_mem)):
                            dis_sum+=geodesic((lat_obs_ave,lon_obs_ave),(lat_mem[l],lon_mem[l])).kilometers
                            dt_sum+=abs((time_mem[l]-time_obs_ave).total_seconds())/3600
            if num==0:
                continue
            num_result[j,i] = num
            dis_result[j,i] = dis_sum/num
            dt_result[j,i]  = dt_sum/num  #hour
    
    
    str_result = np.zeros(num_result.shape,dtype='U50')
    for j in range(num_result.shape[0]):
        for i in range(num_result.shape[1]):
            str_result[j,i] = f"Num:{int(num_result[j,i])}\n Dis_ave:{dis_result[j][i]:.1f}km\n Dt_ave:{dt_result[j][i]:.1f}h"
    
    ''' to excel '''
    df = pd.DataFrame(str_result.T, 
                     columns=[f"{2*j}-{2*(j+1)}天" for j in range(num_result.shape[0])],  # 行标签
                     index=dates_name)
    with pd.ExcelWriter(os.path.join(pic_savepath,f'{ty}_statics.xlsx')) as writer:
        # 将每个 DataFrame 保存到不同的工作表中
        df.to_excel(writer, sheet_name='num_result', index=True)
